import pandas as pd
import numpy as np
import joblib
import datetime
import json
from utils.log import Logger
from features import FeatureEngineering
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score


class LogisticPredictor(object):
    def __init__(self, model_path='../model/logistic_model.pkl', 
                 features_path='../model/logistic_selected_features.json',
                 scaler_path='../model/logistic_scaler.pkl'):
        """初始化逻辑回归预测器"""
        # 配置日志
        predict_log_name = "logistic_predict_" + datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        self.logger = Logger('../log', predict_log_name).get_logger()
        
        # 加载模型、特征列表和标准化器
        self.model = self.load_model(model_path)
        self.selected_features = self.load_features(features_path)
        self.scaler = self.load_scaler(scaler_path)
        
        self.logger.info("逻辑回归预测器初始化完成".center(50, '='))

    def load_model(self, model_path):
        """
        加载已训练好的逻辑回归模型
        :param model_path: 模型路径
        :return: 加载的模型
        """
        self.logger.info("开始加载逻辑回归模型".center(50, '='))
        try:
            model = joblib.load(model_path)
            self.logger.info(f"逻辑回归模型加载成功: {model_path}")
            return model
        except Exception as e:
            self.logger.error(f"逻辑回归模型加载失败: {e}")
            raise e

    def load_features(self, features_path):
        """
        加载训练时选择的特征列表
        :param features_path: 特征列表文件路径
        :return: 特征列表
        """
        self.logger.info("开始加载特征列表".center(50, '='))
        try:
            with open(features_path, 'r') as f:
                selected_features = json.load(f)
            self.logger.info(f"特征列表加载成功: {features_path}")
            self.logger.info(f"特征数量: {len(selected_features)}")
            self.logger.info(f"特征列表: {selected_features}")
            return selected_features
        except Exception as e:
            self.logger.error(f"特征列表加载失败: {e}")
            raise e

    def load_scaler(self, scaler_path):
        """
        加载训练时使用的标准化器
        :param scaler_path: 标准化器路径
        :return: 标准化器
        """
        self.logger.info("开始加载标准化器".center(50, '='))
        try:
            scaler = joblib.load(scaler_path)
            self.logger.info(f"标准化器加载成功: {scaler_path}")
            return scaler
        except Exception as e:
            self.logger.error(f"标准化器加载失败: {e}")
            raise e

    def prepare_data(self, data_path):
        """
        准备预测数据，使用与训练时相同的特征
        :param data_path: 数据路径
        :return: 处理后的特征数据
        """
        self.logger.info("开始准备预测数据".center(50, '='))
        
        # 使用特征工程类进行数据预处理
        fe = FeatureEngineering(data_path)
        x, _ = fe.preprocess()  # 数据预处理
        
        # 检查特征是否匹配
        missing_features = set(self.selected_features) - set(x.columns)
        extra_features = set(x.columns) - set(self.selected_features)
        
        if missing_features:
            self.logger.warning(f"数据缺少以下特征: {missing_features}")
            
        if extra_features:
            self.logger.info(f"数据包含额外特征: {extra_features}")
        
        # 使用训练时选择的特征准备预测数据
        # 对于缺失的特征，用0填充
        x_final = pd.DataFrame(index=x.index, columns=self.selected_features)
        for feature in self.selected_features:
            if feature in x.columns:
                x_final[feature] = x[feature]
            else:
                x_final[feature] = 0  # 用0填充缺失的特征
                
        # 特征标准化
        x_scaled = self.scaler.transform(x_final)
                
        self.logger.info(f"预测数据准备完成，特征数量: {len(self.selected_features)}")
        self.logger.info("预测数据准备完成".center(50, '='))
        
        return x_scaled

    def predict(self, data_path):
        """
        对数据进行预测
        :param data_path: 待预测数据路径
        :return: 预测结果
        """
        self.logger.info("开始进行预测".center(50, '='))
        
        # 准备数据
        x = self.prepare_data(data_path)
        
        # 进行预测
        predictions = self.model.predict(x)
        probabilities = self.model.predict_proba(x)
        
        self.logger.info(f"预测完成，共预测 {len(predictions)} 条数据")
        self.logger.info("预测完成".center(50, '='))
        
        # 返回预测结果和概率
        return predictions, probabilities

    def evaluate_predictions(self, y_true, predictions, probabilities):
        """
        计算并记录常用的评价指标
        :param y_true: 真实标签
        :param predictions: 预测标签
        :param probabilities: 预测概率
        """
        self.logger.info("开始计算评价指标".center(50, '='))
        
        # 计算各种评价指标
        accuracy = accuracy_score(y_true, predictions)
        precision = precision_score(y_true, predictions)
        recall = recall_score(y_true, predictions)
        f1 = f1_score(y_true, predictions)
        auc = roc_auc_score(y_true, probabilities[:, 1])
        
        # 记录评价指标到日志
        self.logger.info("逻辑回归模型预测评价指标:")
        self.logger.info(f"准确率 (Accuracy): {accuracy:.4f}")
        self.logger.info(f"精确率 (Precision): {precision:.4f}")
        self.logger.info(f"召回率 (Recall): {recall:.4f}")
        self.logger.info(f"F1-score: {f1:.4f}")
        self.logger.info(f"AUC: {auc:.4f}")
        
        self.logger.info("评价指标计算完成".center(50, '='))

    def save_predictions(self, predictions, probabilities, output_path='../data/logistic_predictions.csv'):
        """
        保存预测结果
        :param predictions: 预测标签
        :param probabilities: 预测概率
        :param output_path: 输出路径
        """
        self.logger.info("开始保存预测结果".center(50, '='))
        
        # 创建结果DataFrame
        results = pd.DataFrame({
            'Prediction': predictions,
            'Probability_0': probabilities[:, 0],
            'Probability_1': probabilities[:, 1]
        })
        
        # 保存到CSV文件
        results.to_csv(output_path, index=False)
        self.logger.info(f"预测结果已保存至: {output_path}")
        self.logger.info("预测结果保存完成".center(50, '='))


if __name__ == "__main__":
    # 使用示例
    predictor = LogisticPredictor()
    
    # 进行预测（这里假设使用训练数据进行演示）
    predictions, probabilities = predictor.predict('../data/test.csv')
    
    # 保存预测结果
    predictor.save_predictions(predictions, probabilities)
    
    # 如果有真实标签，可以计算评估指标
    try:
        # 尝试读取真实标签
        data = pd.read_csv('../data/test.csv')
        # 需要对标签进行相同的预处理以获取真实标签
        y_true = data['Attrition']
        
        # 计算并记录评估指标
        predictor.evaluate_predictions(y_true, predictions, probabilities)
    except Exception as e:
        predictor.logger.warning(f"无法计算评估指标: {e}")
        predictor.logger.info("如需计算评估指标，请确保数据包含真实标签")